"""
排序代码的封装
"""
import json

import numpy as np
import torch
from torch.optim import Adam
from tqdm import tqdm

from chatbot import config, lib
from chatbot.dnn.sort.dataset import dnnsort_data_loader
from chatbot.dnn.sort.siamese import SiameseNetwork

import torch.nn.functional as F


class DnnSort:
    def __init__(self):
        self.model = SiameseNetwork().to(config.device)
        self.model.load_state_dict(torch.load(config.sort_model_save_path))
        self.model.eval()
        self.qa_dict = json.load(open(config.recall_qa_dict_path, "r"))

    def predict(self, sentence, recall_list):
        input1 = [lib.cut(sentence["cut_by_word"], by_word=True)] * len(recall_list)
        input2 = [self.qa_dict[i]["q_cut_by_word"] for i in recall_list]
        input1 = [config.sort_ws.transform(i, max_len=config.sort_q_max_len) for i in input1]
        input2 = [config.sort_ws.transform(i, max_len=config.sort_q_max_len) for i in input2]
        output = self.model(torch.LongTensor(input1).to(config.device), torch.LongTensor(input2).to(config.device))
        # output = F.softmax(output=self.model(torch.LongTensor(input1).to(config.device), torch.LongTensor(input2).to(config.device)), dim=-1)
        if output.size(0) > 1:
            output = output[:, -1].squeeze(-1).detach().cpu().numpy()
        else:
            output = [output[:, -1].squeeze(-1).detach().cpu().numpy()]
        best_q, best_prod = sorted(list(zip(recall_list, output)), key=lambda x: x[1], reverse=True)[0]

        if best_prod > 0.01:
            return self.qa_dict[best_q]["answer"]
        else:
            # 不相似
            return "无法回答"
